import torch
import torch.nn as nn
import torch.nn.functional as F


class ResUNet(nn.Module):
    def __init__(self, in_channel=1, out_channel=3, training=True):
        super().__init__()

        self.training = training
        self.dorp_rate = 0.2  # 设置dropout的比率为0.2

        # 编码器阶段1
        self.encoder_stage1 = nn.Sequential(
            nn.Conv3d(in_channel, 32, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(32),  # Parametric ReLU激活函数

            nn.Conv3d(32, 32, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(32),  # Parametric ReLU激活函数
        )

        # 编码器阶段2
        self.encoder_stage2 = nn.Sequential(
            nn.Conv3d(64, 64, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(64),  # Parametric ReLU激活函数

            nn.Conv3d(64, 64, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(64),  # Parametric ReLU激活函数

            nn.Conv3d(64, 64, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(64),  # Parametric ReLU激活函数
        )

        # 编码器阶段3
        self.encoder_stage3 = nn.Sequential(
            nn.Conv3d(128, 128, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(128),  # Parametric ReLU激活函数

            nn.Conv3d(128, 128, 3, 1, padding=2, dilation=2),  # 3D卷积层,带孔卷积
            nn.PReLU(128),  # Parametric ReLU激活函数

            nn.Conv3d(128, 128, 3, 1, padding=4, dilation=4),  # 3D卷积层,带孔卷积
            nn.PReLU(128),  # Parametric ReLU激活函数
        )

        # 编码器阶段4
        self.encoder_stage4 = nn.Sequential(
            nn.Conv3d(256, 256, 3, 1, padding=3, dilation=3),  # 3D卷积层,带孔卷积
            nn.PReLU(256),  # Parametric ReLU激活函数

            nn.Conv3d(256, 256, 3, 1, padding=4, dilation=4),  # 3D卷积层,带孔卷积
            nn.PReLU(256),  # Parametric ReLU激活函数

            nn.Conv3d(256, 256, 3, 1, padding=5, dilation=5),  # 3D卷积层,带孔卷积
            nn.PReLU(256),  # Parametric ReLU激活函数
        )

        # 解码器阶段1
        self.decoder_stage1 = nn.Sequential(
            nn.Conv3d(256, 512, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(512),  # Parametric ReLU激活函数

            nn.Conv3d(512, 512, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(512),  # Parametric ReLU激活函数

            nn.Conv3d(512, 512, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(512),  # Parametric ReLU激活函数
        )

        # 解码器阶段2
        self.decoder_stage2 = nn.Sequential(
            nn.Conv3d(256 + 128, 256, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(256),  # Parametric ReLU激活函数

            nn.Conv3d(256, 256, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(256),  # Parametric ReLU激活函数

            nn.Conv3d(256, 256, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(256),  # Parametric ReLU激活函数
        )

        # 解码器阶段3
        self.decoder_stage3 = nn.Sequential(
            nn.Conv3d(128 + 64, 128, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(128),  # Parametric ReLU激活函数

            nn.Conv3d(128, 128, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(128),  # Parametric ReLU激活函数

            nn.Conv3d(128, 128, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(128),  # Parametric ReLU激活函数
        )

        # 解码器阶段4
        self.decoder_stage4 = nn.Sequential(
            nn.Conv3d(64 + 32, 64, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(64),  # Parametric ReLU激活函数

            nn.Conv3d(64, 64, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(64),  # Parametric ReLU激活函数
        )

        # 下采样卷积层
        self.down_conv1 = nn.Sequential(
            nn.Conv3d(32, 64, 2, 2),  # 3D卷积层,[2,2,2]卷积核,步长为2
            nn.PReLU(64)  # Parametric ReLU激活函数
        )

        self.down_conv2 = nn.Sequential(
            nn.Conv3d(64, 128, 2, 2),  # 3D卷积层,[2,2,2]卷积核,步长为2
            nn.PReLU(128)  # Parametric ReLU激活函数
        )

        self.down_conv3 = nn.Sequential(
            nn.Conv3d(128, 256, 2, 2),  # 3D卷积层,[2,2,2]卷积核,步长为2
            nn.PReLU(256)  # Parametric ReLU激活函数
        )

        self.down_conv4 = nn.Sequential(
            nn.Conv3d(256, 512, 3, 1, padding=1),  # 3D卷积层,图像大小不变
            nn.PReLU(512)  # Parametric ReLU激活函数
        )

        # 上采样卷积层
        self.up_conv2 = nn.Sequential(
            nn.ConvTranspose3d(512, 256, 2, 2),  # 3D转置卷积层,[2,2,2]卷积核,步长为2
            nn.PReLU(256)  # Parametric ReLU激活函数
        )

        self.up_conv3 = nn.Sequential(
            nn.ConvTranspose3d(256, 128, 2, 2),  # 3D转置卷积层,[2,2,2]卷积核,步长为2
            nn.PReLU(128)  # Parametric ReLU激活函数
        )

        self.up_conv4 = nn.Sequential(
            nn.ConvTranspose3d(128, 64, 2, 2),  # 3D转置卷积层,[2,2,2]卷积核,步长为2
            nn.PReLU(64)  # Parametric ReLU激活函数
        )

        # 最后大尺度下的映射（256*256），下面的尺度依次递减
        self.map4 = nn.Sequential(
            nn.Conv3d(64, out_channel, 1, 1),  # [1,1,1]卷积核,步长为1
            # 三线性插值上采样
            nn.Upsample(scale_factor=(1, 1, 1), mode='trilinear', align_corners=False)
            # nn.Softmax(dim=1)  # Softmax激活函数
        )

        # 128*128 尺度下的映射
        self.map3 = nn.Sequential(
            nn.Conv3d(128, out_channel, 1, 1),  # [1,1,1]卷积核,步长为1
            # 三线性插值上采样
            nn.Upsample(scale_factor=(2, 2, 2), mode='trilinear', align_corners=False)
            # nn.Softmax(dim=1)  # Softmax激活函数
        )

        # 64*64 尺度下的映射
        self.map2 = nn.Sequential(
            nn.Conv3d(256, out_channel, 1, 1),  # [1,1,1]卷积核,步长为1
            # 三线性插值上采样
            nn.Upsample(scale_factor=(4, 4, 4), mode='trilinear', align_corners=False)
            # nn.Softmax(dim=1)  # Softmax激活函数
        )

        # 32*32 尺度下的映射
        self.map1 = nn.Sequential(
            nn.Conv3d(512, out_channel, 1, 1),  # [1,1,1]卷积核,步长为1
            # 三线性插值上采样
            nn.Upsample(scale_factor=(8, 8, 8), mode='trilinear', align_corners=False)
            # nn.Softmax(dim=1)  # Softmax激活函数
        )

    def forward(self, inputs):
        # 编码器阶段1(in_channel -> 32)
        long_range1 = self.encoder_stage1(inputs) + inputs  # 得到长程残差连接结果
        # 下采样阶段1(32 -> 64)
        short_range1 = self.down_conv1(long_range1)  # 下采样得到短程特征

        # 编码器阶段2(64 -> 64)
        long_range2 = self.encoder_stage2(short_range1) + short_range1  # 得到长程残差连接结果
        long_range2 = F.dropout(long_range2, self.dorp_rate, self.training)  # 使用Dropout进行正则化
        # 下采样阶段2(64 -> 128)
        short_range2 = self.down_conv2(long_range2)  # 下采样得到短程特征

        # 编码器阶段3(128 -> 128)
        long_range3 = self.encoder_stage3(short_range2) + short_range2  # 得到长程残差连接结果
        long_range3 = F.dropout(long_range3, self.dorp_rate, self.training)  # 使用Dropout进行正则化
        # 下采样阶段3(128 -> 256)
        short_range3 = self.down_conv3(long_range3)  # 下采样得到短程特征

        # 编码器阶段4(256 -> 256)
        long_range4 = self.encoder_stage4(short_range3) + short_range3  # 得到长程残差连接结果
        long_range4 = F.dropout(long_range4, self.dorp_rate, self.training)  # 使用Dropout进行正则化
        # 下采样阶段4(256 -> 512)
        short_range4 = self.down_conv4(long_range4)  # 下采样得到短程特征

        # 解码器阶段1输出结果(256 -> 512)
        outputs = self.decoder_stage1(long_range4) + short_range4
        outputs = F.dropout(outputs, self.dorp_rate, self.training)  # 使用Dropout进行正则化

        output1 = self.map1(outputs)  # 经过映射得到最终输出结果output1(512 -> out_channel)

        # 上采样得到短程特征(512 -> 256)
        short_range6 = self.up_conv2(outputs)

        # 解码器阶段2输出结果(256 + 128 -> 256)
        outputs = self.decoder_stage2(torch.cat([short_range6, long_range3], dim=1)) + short_range6
        outputs = F.dropout(outputs, self.dorp_rate, self.training)  # 使用Dropout进行正则化

        output2 = self.map2(outputs)  # 经过映射得到最终输出结果(256 -> out_channel)

        # 上采样得到短程特征(256 -> 128)
        short_range7 = self.up_conv3(outputs)
        # 解码器阶段3输出结果(128 + 64 -> 128)
        outputs = self.decoder_stage3(torch.cat([short_range7, long_range2], dim=1)) + short_range7
        outputs = F.dropout(outputs, self.dorp_rate, self.training)  # 使用Dropout进行正则化

        output3 = self.map3(outputs)  # 经过映射得到最终输出结果(128 -> out_channel)
        # 上采样得到短程特征(128 -> 64)
        short_range8 = self.up_conv4(outputs)
        # 解码器阶段4输出结果(64 + 32 -> 64)
        outputs = self.decoder_stage4(torch.cat([short_range8, long_range1], dim=1)) + short_range8

        output4 = self.map4(outputs)  # 经过映射得到最终输出结果(64 -> out_channel)

        if self.training is True:  # 训练过程返回所有尺度下的输出结果
            return output1, output2, output3, output4
        else:  # 测试过程返回最大尺度下的输出结果
            return output4

